/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.io.gcp.bigquery;



import com.google.api.client.http.HttpRequestInitializer;

import com.google.api.services.bigquery.Bigquery;

import com.google.api.services.bigquery.model.Table;

import com.google.api.services.bigquery.model.TableReference;

import com.google.api.services.bigquery.model.TableRow;

import com.google.api.services.bigquery.model.TableSchema;

import com.google.auth.Credentials;

import com.google.auth.http.HttpCredentialsAdapter;

import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;

import com.bff.gaia.unified.sdk.extensions.gcp.auth.NullCredentialInitializer;

import com.bff.gaia.unified.sdk.extensions.gcp.util.RetryHttpRequestInitializer;

import com.bff.gaia.unified.sdk.extensions.gcp.util.Transport;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.schemas.Schema.FieldType;

import com.bff.gaia.unified.sdk.testing.TestPipeline;

import com.bff.gaia.unified.sdk.testing.TestPipelineOptions;

import com.bff.gaia.unified.sdk.values.Row;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableList;

import org.hamcrest.Matcher;

import org.joda.time.DateTime;

import org.joda.time.Duration;

import org.joda.time.Instant;

import org.joda.time.format.DateTimeFormat;

import org.joda.time.format.DateTimeFormatter;

import org.junit.rules.TestRule;

import org.junit.runner.Description;

import org.junit.runners.model.Statement;



import java.io.IOException;

import java.util.Collections;

import java.util.List;

import java.util.concurrent.ThreadLocalRandom;

import java.util.stream.Collectors;



import static org.joda.time.Seconds.secondsBetween;

import static org.junit.Assert.assertThat;



/**

 * Test rule which creates a new table with specified schema, with randomized name and exposes few

 * APIs to work with it.

 *

 * <p>Deletes the table on test shutdown.

 */

public class TestBigQuery implements TestRule {

  private static final DateTimeFormatter DATETIME_FORMAT =

      DateTimeFormat.forPattern("YYYY_MM_dd_HH_mm_ss_SSS");



  private TestBigQueryOptions pipelineOptions;

  private Schema schema;

  private Table table;

  private BigQueryServices.DatasetService datasetService;



  /**

   * Creates an instance of this rule.

   *

   * <p>Loads GCP configuration from {@link TestPipelineOptions}.

   */

  public static TestBigQuery create(Schema tableSchema) {

    return new TestBigQuery(

        TestPipeline.testingPipelineOptions().as(TestBigQueryOptions.class), tableSchema);

  }



  private TestBigQuery(TestBigQueryOptions pipelineOptions, Schema tableSchema) {

    this.pipelineOptions = pipelineOptions;

    this.schema = tableSchema;

  }



  @Override

  public Statement apply(Statement base, Description description) {

    return new Statement() {

      @Override

      public void evaluate() throws Throwable {

        if (TestBigQuery.this.datasetService != null) {

          throw new AssertionError(

              "BigQuery test was not shutdown previously. "

                  + "Table is'"

                  + table

                  + "'. "

                  + "Current test: "

                  + description.getDisplayName());

        }



        try {

          initializeBigQuery(description);

          base.evaluate();

        } finally {

          tearDown();

        }

      }

    };

  }



  private void initializeBigQuery(Description description)

      throws IOException, InterruptedException {



    this.datasetService = new BigQueryServicesImpl().getDatasetService(pipelineOptions);

    this.table = createTable(description);

  }



  private Table createTable(Description description) throws IOException, InterruptedException {

    TableReference tableReference =

        new TableReference()

            .setProjectId(pipelineOptions.getProject())

            .setDatasetId(pipelineOptions.getTargetDataset())

            .setTableId(createRandomizedName(description));



    table =

        new Table()

            .setTableReference(tableReference)

            .setSchema(BigQueryUtils.toTableSchema(schema))

            .setDescription(

                "Table created for "

                    + description.getDisplayName()

                    + " by TestBigQueryRule. "

                    + "Should be automatically cleaned up after test completion.");



    if (datasetService.getTable(tableReference) != null) {

      throw new IllegalStateException(

          "Table '"

              + tableReference

              + "' already exists. "

              + "It should have been cleaned up by the test rule.");

    }



    datasetService.createTable(table);

    return table;

  }



  private void tearDown() throws IOException, InterruptedException {

    if (this.datasetService == null) {

      return;

    }



    try {

      if (table != null) {

        datasetService.deleteTable(table.getTableReference());

      }

    } finally {

      datasetService = null;

      table = null;

    }

  }



  static String createRandomizedName(Description description) throws IOException {

    StringBuilder topicName = new StringBuilder();



    if (description.getClassName() != null) {

      try {

        topicName.append(Class.forName(description.getClassName()).getSimpleName()).append("_");

      } catch (ClassNotFoundException e) {

        throw new RuntimeException(e);

      }

    }



    if (description.getMethodName() != null) {

      topicName.append(description.getMethodName()).append("_");

    }



    DATETIME_FORMAT.printTo(topicName, Instant.now());



    return topicName.toString()

        + "_"

        + String.valueOf(Math.abs(ThreadLocalRandom.current().nextLong()));

  }



  public String tableSpec() {

    return String.format(

        "%s:%s.%s",

        table.getTableReference().getProjectId(),

        table.getTableReference().getDatasetId(),

        table.getTableReference().getTableId());

  }



  /**

   * Loads rows from BigQuery into {@link Row Rows} with given {@link Schema}.

   *

   * <p>Current implementation only supports flat {@link Row Rows} and target {@link Schema Schemas}

   * with {@link FieldType#STRING} fields only.

   */

  public List<Row> getFlatJsonRows(Schema rowSchema) {

    Bigquery bq = newBigQueryClient(pipelineOptions);

    return bqRowsToUnifiedRows(getSchema(bq), getTableRows(bq), rowSchema);

  }



  public RowsAssertion assertThatAllRows(Schema rowSchema) {

    return matcher -> duration -> pollAndAssert(rowSchema, matcher, duration);

  }



  private void pollAndAssert(

	  Schema rowSchema, Matcher<Iterable<? extends Row>> matcher, Duration duration) {



    DateTime start = DateTime.now();

    while (true) {

      try {

        assertThat(getFlatJsonRows(rowSchema), matcher);

        break;

      } catch (AssertionError assertionError) {

        if (secondsBetween(start, DateTime.now()).isGreaterThan(duration.toStandardSeconds())) {

          throw assertionError;

        }

        sleep(15_000);

      }

    }

  }



  private List<Row> bqRowsToUnifiedRows(

      TableSchema bqSchema, List<TableRow> bqRows, Schema rowSchema) {

    if (bqRows == null) {

      return Collections.emptyList();

    }



    return bqRows.stream()

        .map(bqRow -> BigQueryUtils.toUnifiedRow(rowSchema, bqSchema, bqRow))

        .collect(Collectors.toList());

  }



  private TableSchema getSchema(Bigquery bq) {

    try {

      return bq.tables()

          .get(

              pipelineOptions.getProject(),

              pipelineOptions.getTargetDataset(),

              table.getTableReference().getTableId())

          .execute()

          .getSchema();

    } catch (IOException e) {

      throw new RuntimeException(e);

    }

  }



  private List<TableRow> getTableRows(Bigquery bq) {

    try {

      return bq.tabledata()

          .list(

              pipelineOptions.getProject(),

              pipelineOptions.getTargetDataset(),

              table.getTableReference().getTableId())

          .execute()

          .getRows();

    } catch (IOException e) {

      throw new RuntimeException(e);

    }

  }



  private static Bigquery newBigQueryClient(BigQueryOptions options) {

    return new Bigquery.Builder(

            Transport.getTransport(),

            Transport.getJsonFactory(),

            chainHttpRequestInitializer(

                options.getGcpCredential(),

                // Do not log 404. It clutters the output and is possibly even required by the

                // caller.

                new RetryHttpRequestInitializer(ImmutableList.of(404))))

        .setApplicationName(options.getAppName())

        .setGoogleClientRequestInitializer(options.getGoogleApiTrace())

        .build();

  }



  private static HttpRequestInitializer chainHttpRequestInitializer(

      Credentials credential, HttpRequestInitializer httpRequestInitializer) {

    if (credential == null) {

      return new ChainingHttpRequestInitializer(

          new NullCredentialInitializer(), httpRequestInitializer);

    } else {

      return new ChainingHttpRequestInitializer(

          new HttpCredentialsAdapter(credential), httpRequestInitializer);

    }

  }



  private void sleep(long l) {

    try {

      Thread.sleep(l);

    } catch (InterruptedException e) {

      throw new RuntimeException(e);

    }

  }



  /** Interface for creating a polling eventual assertion. */

  public interface RowsAssertion {

    PollingAssertion eventually(Matcher<Iterable<? extends Row>> matcher);

  }



  /** Interface to implement a polling assertion. */

  public interface PollingAssertion {

    void pollFor(Duration duration);

  }

}